package collector

import (
	"context"
	"os"
	"sort"
	"time"

	"sysinfo/internal/types"

	psnet "github.com/shirou/gopsutil/v3/net"
	"github.com/shirou/gopsutil/v3/process"
)

// CollectPorts 收集各平台的监听端口占用信息（TCP/UDP）
func CollectPorts() (*types.PortsData, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	tcpConns, _ := psnet.ConnectionsWithContext(ctx, "tcp")
	udpConns, _ := psnet.ConnectionsWithContext(ctx, "udp")

	entries := make([]types.PortEntry, 0, len(tcpConns)+len(udpConns))
	total := 0
	procNameCache := map[int32]string{}

	// 缓存并获取进程名
	getProcName := func(pid int32) string {
		if pid <= 0 {
			return ""
		}
		if name, ok := procNameCache[pid]; ok {
			return name
		}
		p, err := process.NewProcess(pid)
		if err != nil {
			procNameCache[pid] = ""
			return ""
		}
		name, err := p.Name()
		if err != nil {
			procNameCache[pid] = ""
			return ""
		}
		procNameCache[pid] = name
		return name
	}

	// TCP: 仅统计 LISTEN 状态的端口
	for _, c := range tcpConns {
		if c.Status == "LISTEN" || c.Status == "LISTENING" {
			total++
			entries = append(entries, types.PortEntry{
				Protocol:    "tcp",
				LocalIP:     c.Laddr.IP,
				LocalPort:   uint32(c.Laddr.Port),
				PID:         c.Pid,
				ProcessName: getProcName(c.Pid),
			})
		}
	}

	// UDP: 没有 LISTEN 状态，任何存在本地端口的条目视为监听
	for _, c := range udpConns {
		if c.Laddr.Port != 0 {
			total++
			entries = append(entries, types.PortEntry{
				Protocol:    "udp",
				LocalIP:     c.Laddr.IP,
				LocalPort:   uint32(c.Laddr.Port),
				PID:         c.Pid,
				ProcessName: getProcName(c.Pid),
			})
		}
	}

	// 稳定排序：协议 -> 端口 -> IP
	sort.Slice(entries, func(i, j int) bool {
		if entries[i].Protocol != entries[j].Protocol {
			return entries[i].Protocol < entries[j].Protocol
		}
		if entries[i].LocalPort != entries[j].LocalPort {
			return entries[i].LocalPort < entries[j].LocalPort
		}
		return entries[i].LocalIP < entries[j].LocalIP
	})

	return &types.PortsData{
		TotalListening: total,
		Entries:        entries,
	}, nil
}

// KillProcessesByPorts 根据给定端口列表，终止正在监听这些端口的进程（跨平台）
// Windows、Linux、macOS 均通过标准库 os.Process.Kill 进行终止。
// 注意：可能需要管理员/特权权限，非当前用户进程通常无法被终止。
func KillProcessesByPorts(ports []int) (killed []int32, failed map[int32]error, err error) {
	if len(ports) == 0 {
		return nil, nil, nil
	}

	// 端口集合去重
	portSet := make(map[uint32]struct{}, len(ports))
	for _, p := range ports {
		if p > 0 {
			portSet[uint32(p)] = struct{}{}
		}
	}

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	tcpConns, _ := psnet.ConnectionsWithContext(ctx, "tcp")
	udpConns, _ := psnet.ConnectionsWithContext(ctx, "udp")

	// 收集需要终止的 PID
	pidSet := make(map[int32]struct{})
	for _, c := range tcpConns {
		if _, ok := portSet[uint32(c.Laddr.Port)]; ok {
			if c.Status == "LISTEN" || c.Status == "LISTENING" {
				if c.Pid > 0 {
					pidSet[c.Pid] = struct{}{}
				}
			}
		}
	}
	for _, c := range udpConns {
		if _, ok := portSet[uint32(c.Laddr.Port)]; ok {
			if c.Pid > 0 {
				pidSet[c.Pid] = struct{}{}
			}
		}
	}

	failed = make(map[int32]error)

	// 逐个终止进程
	for pid := range pidSet {
		proc, e := os.FindProcess(int(pid))
		if e != nil {
			failed[pid] = e
			continue
		}
		if e = proc.Kill(); e != nil {
			failed[pid] = e
			continue
		}
		killed = append(killed, pid)
	}

	// 排序以稳定输出
	sort.Slice(killed, func(i, j int) bool { return killed[i] < killed[j] })

	return killed, failed, nil
}
